#!/usr/bin/python3 -u
# Add content on top of a commit, handling SELinux labeling etc.
# This is a more flexible (and correct) replacement for the CI flow developed
# in https://github.com/smarterclayton/origin/blob/4de957b019aee56931b1a29af148cf64865a969b/images/os/Dockerfile
# Down the line we want to lower this into rpm-ostree (and support RPMs too).
# This script though should be very convenient for the common case of
# overlaying some new binaries into /usr/bin that don't need any scripts
# run and also don't need e.g. to regenerate the initramfs that would
# happen in a full `rpm-ostree compose tree`.
# Also related to the initramfs, this script explicitly doesn't
# use or require any container features, so it will work inside
# an podman/Kube container without special privileges.

import argparse
import gi
import os
import tempfile

gi.require_version('OSTree', '1.0')
gi.require_version('Json', '1.0')
from gi.repository import Gio, OSTree

# This seems to not be exposed by Python which has its own wrappers,
# It's also not exposed by GLib.  This is a Linux specific value
# but that's totally fine, OSTree only works there too.
AT_FDCWD = -100

parser = argparse.ArgumentParser()
parser.add_argument("--repo", help="repo", required=True)
parser.add_argument("--rev", help="Revision to override")
# parser.add_argument("--add-rpm", help="Unpack literal RPM content (does not currently update rpm DB, replace older RPM, or run any scripts)")
parser.add_argument("--add-tree", help="Add local filesystem tree", default=[], action="append")
parser.add_argument("--output-ref", help="Output ref (if unset, will just write new commit)")
args = parser.parse_args()

repo = OSTree.Repo.new(Gio.File.new_for_path(args.repo))
repo.open(None)

if args.rev is not None:
    rev = args.rev
else:
    [_, refs] = repo.list_refs(None, None)
    nrefs = len(refs)
    if nrefs > 1:
        raise SystemExit(f"No rev specified and repo has {nrefs} refs")
    rev = refs[0]
[_, root, rev] = repo.read_commit(rev, None)
[_, base_commit, _] = repo.load_commit(rev)
# See https://github.com/ostreedev/ostree/pull/1643
base_contents_checksum = OSTree.checksum_from_bytes_v(base_commit.get_child_value(6))
base_meta_checksum = OSTree.checksum_from_bytes_v(base_commit.get_child_value(7))
mtree = OSTree.MutableTree.new_from_checksum(repo, base_contents_checksum, base_meta_checksum)

tmpd = tempfile.TemporaryDirectory(dir=f"{args.repo}/tmp/", prefix="dev-overlay")


def add_commit_filter(repo, path, finfo):
    # Canonicalize uid/gid to 0
    finfo.set_attribute_uint32("unix::uid", 0)
    finfo.set_attribute_uint32("unix::gid", 0)
    return OSTree.RepoCommitFilterResult.ALLOW


add_modifier = OSTree.RepoCommitModifier.new(OSTree.RepoCommitModifierFlags.DEVINO_CANONICAL |
                                             OSTree.RepoCommitModifierFlags.SKIP_XATTRS, add_commit_filter)
if root.get_child("usr/etc/selinux"):
    opts = OSTree.RepoCheckoutAtOptions()
    opts.mode = OSTree.RepoCheckoutMode.USER
    opts.subpath = "/usr/etc/selinux"
    dest = tmpd.name + "/" + opts.subpath
    os.makedirs(os.path.dirname(dest))
    repo.checkout_at(opts, AT_FDCWD, dest, rev, None)
    add_modifier.set_sepolicy(OSTree.SePolicy.new(Gio.File.new_for_path(tmpd.name)))
for d in args.add_tree:
    repo.write_dfd_to_mtree(AT_FDCWD, d, mtree, add_modifier, None)

[_, dir_tree] = repo.write_mtree(mtree, None)
[_, new_commit] = repo.write_commit(None, None, None, None, dir_tree, None)
if args.output_ref is not None:
    repo.set_ref_immediate(None, args.output_ref, new_commit)
    print(f"Wrote {args.output_ref} => {new_commit}")
else:
    print(f"Wrote {new_commit}")
